# -*- coding: utf-8 -*-

import os
import cv2
import numpy as np
from PIL import Image
import os

os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"

import torch
import torch.nn as nn
from torchvision import transforms

from models.net_builder import net_builder


def segmenting_DR(img_path, gpu_ids=""):
    # 是否使用cuda
    if gpu_ids != "" and torch.cuda.is_available():
        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # 加载模型
    model = net_builder('camwnet', None, False)
    model = nn.DataParallel(model).to(device)
    checkpoint = torch.load("camw.pkl", map_location=device)
    model.load_state_dict(checkpoint["net"])  # ##
    model.eval()  # ##

    # 读入并预处理图像
    resize_img = transforms.Resize((512, 512), Image.BILINEAR)
    img = Image.open(img_path)
    w, h = img.size
    img = resize_img(img)
    img = np.array(img).astype(np.uint8)
    img = img.transpose(2, 0, 1) / 255.0
    x = torch.from_numpy(img.copy()).float().unsqueeze(0).to(device)

    # 模型预测
    # {"background": 0, "hemorrhages出血": 1, "hard_exudates硬性渗出": 2, "microaneurysms微动脉瘤": 3, "disc": 4, "soft_exudates软性渗出斑": 5}
    with torch.no_grad():
        if device == torch.device("cuda"):
            y = torch.exp(model(x)[0]).data.squeeze().cpu().numpy() * 255.0
        else:
            y = torch.exp(model.module(x)[0]).data.squeeze().cpu().numpy() * 255.0

    predict_list = []
    for index, item in enumerate(range(y.shape[0])):
        predict = Image.fromarray(y[index].squeeze().astype(np.uint8), mode='L')  # 灰度模式存图，h * w无通道
        predict_resize = np.array(predict.resize([w, h]))
        predict_list.append(predict_resize)

    prob_arr = np.array(predict_list) / 255.0
    prob_arr /= np.sum(prob_arr, axis=0)
    pred_arr = np.zeros_like(prob_arr).astype(np.uint8)
    pred_label_arr = np.argmax(prob_arr, axis=0)

    for j in range(pred_arr.shape[0]):
        pred_arr[j, :, :][pred_label_arr == j] = 255

    return pred_arr  # [6, H, W]


def showing_seg(img_path, pred_arr):
    img_arr = cv2.imread(img_path)
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    for j in range(1, pred_arr.shape[0]):
        edge_arr = cv2.dilate(cv2.Canny(pred_arr[j, :, :], 50, 150), kernel, iterations=1)
        edge_arr = cv2.cvtColor(edge_arr, cv2.COLOR_GRAY2BGR)
        img_arr *= (1 - edge_arr // 255)
        if j // 4 == 0:
            edge_arr[:, :, 0] = 0
        if j % 4 < 2:
            edge_arr[:, :, 1] = 0
        if j % 2 == 0:
            edge_arr[:, :, -1] = 0
        img_arr += edge_arr
    return img_arr  # [H, W, 3]


def fsegmenting_DR(img_arr, gpu_ids=""):
    # 是否使用cuda
    if gpu_ids != "" and torch.cuda.is_available():
        os.environ["CUDA_VISIBLE_DEVICES"] = gpu_ids
        device = torch.device("cuda")
    else:
        device = torch.device("cpu")

    # 加载模型
    model = net_builder('camwnet', None, False)
    model = nn.DataParallel(model).to(device)
    checkpoint = torch.load("camw.pkl", map_location=device)
    model.load_state_dict(checkpoint["net"])  # ##
    model.eval()  # ##

    # 读入并预处理图像
    resize_img = transforms.Resize((512, 512), Image.BILINEAR)
    img = Image.fromarray(np.uint8(img_arr))
    w, h = img.size
    img = resize_img(img)
    img = np.array(img).astype(np.uint8)
    img = img.transpose(2, 0, 1) / 255.0
    x = torch.from_numpy(img.copy()).float().unsqueeze(0).to(device)

    # 模型预测
    # {"background": 0, "hemorrhages出血": 1, "hard_exudates硬性渗出": 2, "microaneurysms微动脉瘤": 3, "disc": 4, "soft_exudates软性渗出斑": 5}
    with torch.no_grad():
        if device == torch.device("cuda"):
            y = torch.exp(model(x)[0]).data.squeeze().cpu().numpy() * 255.0
        else:
            y = torch.exp(model.module(x)[0]).data.squeeze().cpu().numpy() * 255.0

    predict_list = []
    for index, item in enumerate(range(y.shape[0])):
        predict = Image.fromarray(y[index].squeeze().astype(np.uint8), mode='L')  # 灰度模式存图，h * w无通道
        predict_resize = np.array(predict.resize([w, h]))
        predict_list.append(predict_resize)

    prob_arr = np.array(predict_list) / 255.0
    prob_arr /= np.sum(prob_arr, axis=0)
    pred_arr = np.zeros_like(prob_arr).astype(np.uint8)
    pred_label_arr = np.argmax(prob_arr, axis=0)

    for j in range(pred_arr.shape[0]):
        pred_arr[j, :, :][pred_label_arr == j] = 255

    return pred_arr  # [6, H, W]


def fshowing_seg(img_arr, pred_arr):
    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    for j in range(1, pred_arr.shape[0]):
        edge_arr = cv2.dilate(cv2.Canny(pred_arr[j, :, :], 50, 150), kernel, iterations=1)
        edge_arr = cv2.cvtColor(edge_arr, cv2.COLOR_GRAY2BGR)
        img_arr *= (1 - edge_arr // 255)
        if j // 4 == 0:
            edge_arr[:, :, 0] = 0
        if j % 4 < 2:
            edge_arr[:, :, 1] = 0
        if j % 2 == 0:
            edge_arr[:, :, -1] = 0
        img_arr += edge_arr
    return img_arr  # [H, W, 3]


if __name__ == '__main__':
    for i in range(1, 6):
        s = str(i)
        print(s)
        pred_arr = segmenting_DR('pic/' + s + '.jpeg')
        img_arr = showing_seg('pic/' + s + '.jpeg', pred_arr)
        cv2.imwrite('re/vis_' + s + '.png', img_arr)
        predict_array = np.zeros_like(pred_arr).astype(np.uint8)
        pred_flag_arr = np.argmax(pred_arr, axis=0)

        for j in range(predict_array.shape[0]):
            predict_array[j, :, :][pred_flag_arr == j] = 255
            cv2.imwrite(str(j) + '_' + s + '.png', predict_array[j, :, :])

        pred_color_arr = np.zeros((pred_flag_arr.shape[0], pred_flag_arr.shape[1], 3)).astype(np.uint8)
        pred_color_arr[:, :, 0][pred_flag_arr // 4 > 0] = 255
        pred_color_arr[:, :, 1][pred_flag_arr % 4 > 1] = 255
        pred_color_arr[:, :, -1][pred_flag_arr % 2 > 0] = 255
        cv2.imwrite('all_' + s + '.png', pred_color_arr)
        print(s + '.jpeg done. ')
